#pragma once

#include <onnxruntime_cxx_api.h>

#include <opencv2/dnn/dnn.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

#include <chrono>
#include <cmath>
#include <exception>
#include <fstream>
#include <iostream>
#include <limits>
#include <numeric>
#include <string>
#include <vector>
#include <math.h>

#include "BaseDetector.hxx"

#define CONF_THRESHOLD 0.15f
#define IOU_THRESHOLD 0.25f

class YOLODetector :
    public BaseDetector
{
public:
    YOLODetector();

    void initialize();  

    void detect(const cv::Mat& image, std::vector<DetectedObject>& retObjs);

private:
    void yolo_nms(const std::vector<cv::Rect>& srcRects, std::vector<float>& confidences, std::vector<cv::Rect>& resRects, std::vector<int>& resIndexs, float thresh = IOU_THRESHOLD);

private:
    Ort::Session* _ortSession;
    Ort::Env* _ortEnv;
    Ort::SessionOptions* _ortSessionOptions;

    std::vector<int64_t> _inputDims;
    std::vector<int64_t> _outputDims;

    char _inputName[256];
    char _outputName[256];

    bool _isInitialized;
};

